package com.foreveross.bsl.common.utils.balance.cxf;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
import org.apache.cxf.endpoint.AbstractConduitSelector;
import org.apache.cxf.endpoint.Client;
import org.apache.cxf.endpoint.Endpoint;
import org.apache.cxf.endpoint.Retryable;
import org.apache.cxf.helpers.CastUtils;
import org.apache.cxf.message.Exchange;
import org.apache.cxf.message.Message;
import org.apache.cxf.service.model.BindingOperationInfo;
import org.apache.cxf.transport.Conduit;

import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.logging.Logger;

/**
 * Implements a target selection strategy based on load balance to an
 * alternate target endpoint.
 *
 * @author Vector Ho  Date:2014-10-17 23:26
 * @see org.apache.cxf.clustering.FailoverTargetSelector
 */
public class LoadBalanceTargetSelector extends AbstractConduitSelector {

    private static final Log logger = LogFactory.getLog(LoadBalanceTargetSelector.class);

    private Map<InvocationKey, InvocationContext> inProgress;
    private LoadBalanceStrategy loadBalanceStrategy;

    /**
     * Normal constructor.
     */
    public LoadBalanceTargetSelector() {
        this(null);
    }

    /**
     * Constructor, allowing a specific conduit to override normal selection.
     *
     * @param c specific conduit
     */
    public LoadBalanceTargetSelector(Conduit c) {
        super(c);
        inProgress = new ConcurrentHashMap<InvocationKey, InvocationContext>();
    }

    /**
     * Called prior to the interceptor chain being traversed.
     *
     * @param message the current Message
     */
    public synchronized void prepare(Message message) {
        Exchange exchange = message.getExchange();
        InvocationKey key = new InvocationKey(exchange);
        if (!inProgress.containsKey(key)) {

            boolean existsEndpoint = false;
            //check current endpoint is not null
            Endpoint theEndpoint = exchange.get(Endpoint.class);
            if (theEndpoint.getEndpointInfo().getAddress() != null) {
                existsEndpoint = true;
            }

            Endpoint nextEndpoint;
            if (getStrategy().isAlwaysChangeEndpoint() || !existsEndpoint) {
                //get a endpoint and set to current endpoint
                Endpoint loadBalanceTarget = getLoadBalanceTarget(exchange);
                if (loadBalanceTarget != null) {
                    logger.info("switch to next target: " + loadBalanceTarget.getEndpointInfo().getAddress());
                    setEndpoint(loadBalanceTarget);

                    //update exchange.org.apache.cxf.message.Message.ENDPOINT_ADDRESS
                    message.put(Message.ENDPOINT_ADDRESS, loadBalanceTarget.getEndpointInfo().getAddress());
                }

                nextEndpoint = loadBalanceTarget;
            } else {
                //use current endpoint
                nextEndpoint = theEndpoint;
            }

            BindingOperationInfo bindingOperationInfo = exchange.getBindingOperationInfo();
            Object[] params = message.getContent(List.class).toArray();
            Map<String, Object> context =  CastUtils.cast((Map) message.get(Message.INVOCATION_CONTEXT));

            InvocationContext invocation =
                    new InvocationContext(nextEndpoint,
                            bindingOperationInfo,
                            params,
                            context);
            inProgress.put(key, invocation);
        }
    }

    /**
     * Called when a Conduit is actually required.
     *
     * @param message message
     * @return the Conduit to use for mediation of the message
     */
    public Conduit selectConduit(Message message) {
        return getSelectedConduit(message);
    }

    /**
     * Called on completion of the MEP for which the Condit was required.
     *
     * @param exchange represents the completed MEP
     */
    public void complete(Exchange exchange) {
        InvocationKey key = new InvocationKey(exchange);
        InvocationContext invocation = null;
        synchronized (this) {
            invocation = inProgress.get(key);
        }
        boolean failover = false;
        if (requiresFailover(exchange)) {
            Endpoint failoverTarget = getFailoverTarget(exchange, invocation);
            if (failoverTarget != null) {
                logger.info("switch to failover target: " + failoverTarget.getEndpointInfo().getAddress());

                setEndpoint(failoverTarget);
                
                try {
                    if (exchange.getInMessage() != null) {
                        Conduit c = exchange.getOutMessage().get(Conduit.class);
                        if (c == null) {
                            getSelectedConduit(exchange.getInMessage()).close(exchange.getInMessage());
                        } else {
                            c.close(exchange.getInMessage());
                        }
                    }
                } catch (IOException e) {
                    //IGNORE
                }
                // 修改上面关闭程序
                // selectedConduit.close();
                // selectedConduit = null;
                
                Exception prevExchangeFault = (Exception) exchange.remove(Exception.class.getName());
                Message outMessage = exchange.getOutMessage();
                Exception prevMessageFault = outMessage.getContent(Exception.class);
                outMessage.setContent(Exception.class, null);
                overrideAddressProperty(invocation.getContext());
                Retryable retry = exchange.get(Retryable.class);
                exchange.clear();
                if (retry != null) {
                    try {
                        failover = true;
                        retry.invoke(invocation.getBindingOperationInfo(),
                                invocation.getParams(),
                                invocation.getContext(),
                                exchange);
                    }
                    catch (Exception e) {
                        if (exchange.get(Exception.class) != null) {
                            exchange.put(Exception.class, prevExchangeFault);
                        }
                        if (outMessage.getContent(Exception.class) != null) {
                            outMessage.setContent(Exception.class,
                                    prevMessageFault);
                        }
                    }
                }
            } else {
                setEndpoint(invocation.retrieveOriginalEndpoint(endpoint));
            }
        }
        if (!failover) {
            logger.info("Failover not required");
            synchronized (this) {
                inProgress.remove(key);
            }
            super.complete(exchange);
        }
    }

    /**
     * @param strategy the FailoverStrategy to use
     */
    public synchronized void setStrategy(LoadBalanceStrategy strategy) {
        logger.info("set strategy:" + strategy);
        loadBalanceStrategy = strategy;
    }

    /**
     * @return strategy the FailoverStrategy to use
     */
    public synchronized LoadBalanceStrategy getStrategy() {
        if (loadBalanceStrategy == null) {
            loadBalanceStrategy = new RandomLoadBalanceStrategy();
            logger.info("using strategy:" + loadBalanceStrategy);
        }
        return loadBalanceStrategy;
    }

    /**
     * @return the logger to use
     */
    protected Logger getLogger() {
        throw new RuntimeException("not support now");
    }

    /**
     * Check if the exchange is suitable for a failover.
     *
     * @param exchange the current Exchange
     * @return boolean true if a failover should be attempted
     */
    private boolean requiresFailover(Exchange exchange) {
        Message outMessage = exchange.getOutMessage();
        Exception ex = outMessage.get(Exception.class) != null
                ? outMessage.get(Exception.class)
                : exchange.get(Exception.class);
        logger.debug(
                "check last invoke failed: " + (ex != null));
        Throwable curr = ex;
        boolean failover = false;
        while (curr != null) {
            failover = curr instanceof java.io.IOException;
            curr = curr.getCause();
        }
        logger.info("check failure in transport");
        return failover;
    }

    /**
     * Get the loadbalance target endpoint, if a suitable one is available.
     *
     * @param exchange the current Exchange
     * @return a loadbalance endpoint if one is available
     */
    private Endpoint getLoadBalanceTarget(Exchange exchange) {
        List<String> alternateAddresses = null;

        alternateAddresses = getStrategy().getAlternateAddresses(exchange);

        Endpoint nextTarget = null;
        if (alternateAddresses != null) {
            String alternateAddress =
                    getStrategy().selectAlternateAddress(alternateAddresses);
            if (alternateAddress != null) {
                nextTarget = getEndpoint();
                nextTarget.getEndpointInfo().setAddress(alternateAddress);
            }
        } else {
            nextTarget = getStrategy().selectAlternateEndpoint(
                    getStrategy().getAlternateEndpoints(exchange));
        }

        return nextTarget;
    }

    /**
     * Get the failover target endpoint, if a suitable one is available.
     *
     * @param exchange   the current Exchange
     * @param invocation the current InvocationContext
     * @return a failover endpoint if one is available
     */
    private Endpoint getFailoverTarget(Exchange exchange,
                                       InvocationContext invocation) {

        List<String> alternateAddresses = null;
        if (!invocation.hasAlternates()) {
            // no previous failover attempt on this invocation
            //
            alternateAddresses =
                    getStrategy().getAlternateAddresses(exchange);
            if (alternateAddresses != null) {
                invocation.setAlternateAddresses(alternateAddresses);
            } else {
                invocation.setAlternateEndpoints(
                        getStrategy().getAlternateEndpoints(exchange));
            }
        } else {
            alternateAddresses = invocation.getAlternateAddresses();
        }

        Endpoint failoverTarget = null;
        if (alternateAddresses != null) {

            String currentAddress = getEndpoint().getEndpointInfo().getAddress();

            //failover should remove current endpoint first, then get next
            if (getStrategy().isRemoveFailedEndpoint()) {
                logger.warn("remove current failed address: " + currentAddress);
                //remove for client, not for current invocation
                getStrategy().removeAlternateAddress(currentAddress);
            }

            //remove for current invocation
            alternateAddresses.remove(currentAddress);

            String alternateAddress =
                    getStrategy().selectAlternateAddress(alternateAddresses);
            if (alternateAddress != null) {
                // re-use current endpoint
                //
                failoverTarget = getEndpoint();

                failoverTarget.getEndpointInfo().setAddress(alternateAddress);
            }
        } else {
            if (getStrategy().isRemoveFailedEndpoint()) {
                //@WARN remove failed endpoint first
                //don't know how to remove
            }

            //remove for current invocation
            invocation.removeCurrentEndpoint();

            failoverTarget = getStrategy().selectAlternateEndpoint(
                    invocation.getAlternateEndpoints());
        }
        return failoverTarget;
    }

    /**
     * Override the ENDPOINT_ADDRESS property in the request context
     *
     * @param context the request context
     */
    private void overrideAddressProperty(Map<String, Object> context) {
        Map<String, Object> requestContext =
                CastUtils.cast((Map) context.get(Client.REQUEST_CONTEXT));
        if (requestContext != null) {
            requestContext.put(Message.ENDPOINT_ADDRESS,
                    getEndpoint().getEndpointInfo().getAddress());
            requestContext.put("javax.xml.ws.service.endpoint.address",
                    getEndpoint().getEndpointInfo().getAddress());
        }
    }

    /**
     * Used to wrap an Exchange for usage as a Map key. The raw Exchange
     * is not a suitable key type, as the hashCode is computed from its
     * current contents, which may obvioulsy change over the lifetime of
     * an invocation.
     */
    private static class InvocationKey {
        private Exchange exchange;

        InvocationKey(Exchange ex) {
            exchange = ex;
        }

        @Override
        public int hashCode() {
            return System.identityHashCode(exchange);
        }

        @Override
        public boolean equals(Object o) {
            return o instanceof InvocationKey
                    && exchange == ((InvocationKey) o).exchange;
        }
    }


    /**
     * Records the context of an invocation.
     */
    private class InvocationContext {
        private Endpoint originalEndpoint;
        private String originalAddress;
        private BindingOperationInfo bindingOperationInfo;
        private Object[] params;
        private Map<String, Object> context;
        private List<Endpoint> alternateEndpoints;
        private List<String> alternateAddresses;

        InvocationContext(Endpoint endpoint,
                          BindingOperationInfo boi,
                          Object[] prms,
                          Map<String, Object> ctx) {
            originalEndpoint = endpoint;
            originalAddress = endpoint.getEndpointInfo().getAddress();
            bindingOperationInfo = boi;
            params = prms;
            context = ctx;
        }

        Endpoint retrieveOriginalEndpoint(Endpoint endpoint) {
            if (endpoint != originalEndpoint) {
                logger.info("revert to original target: " + endpoint.getEndpointInfo().getName());
            }
            if (!endpoint.getEndpointInfo().getAddress().equals(originalAddress)) {
                endpoint.getEndpointInfo().setAddress(originalAddress);
                logger.info("revert to original address: " + endpoint.getEndpointInfo().getAddress());
            }
            return originalEndpoint;
        }

        BindingOperationInfo getBindingOperationInfo() {
            return bindingOperationInfo;
        }

        Object[] getParams() {
            return params;
        }

        Map<String, Object> getContext() {
            return context;
        }

        List<Endpoint> getAlternateEndpoints() {
            return alternateEndpoints;
        }

        List<String> getAlternateAddresses() {
            return alternateAddresses;
        }

        void setAlternateEndpoints(List<Endpoint> alternates) {
            alternateEndpoints = alternates;
        }

        void setAlternateAddresses(List<String> alternates) {
            alternateAddresses = alternates;
        }

        void removeCurrentEndpoint() {
            if (alternateEndpoints != null) {
                alternateEndpoints.remove(originalEndpoint);
            }
        }

        boolean hasAlternates() {
            return !(alternateEndpoints == null && alternateAddresses == null);
        }
    }
}
